"""
Database abstraction layer for CodeMCP storage.

Provides async database operations using SQLAlchemy with support for
SQLite, PostgreSQL, and other databases.
"""

import asyncio
import os
from contextlib import asynccontextmanager
from typing import Dict, Any, List, Optional, Union, Type, TypeVar
from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import (
    create_async_engine, AsyncSession, async_sessionmaker
)
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from pathlib import Path
import json
from datetime import datetime, timedelta

from .models import (
    Base, CallTreeModel, CallTreeNodeModel, AnalysisResultModel,
    FileChangeModel, CacheEntryModel
)
from ..core.config import Config
from ..core.error_handler import (
    CodeMCPError, log_info, log_warning, log_debug, handle_error
)

T = TypeVar('T', bound=Base)


class DatabaseError(CodeMCPError):
    """Database-specific errors."""
    
    def __init__(self, message: str, operation: str = None, **kwargs):
        super().__init__(message, "DATABASE_ERROR", **kwargs)
        if operation:
            self.details["operation"] = operation


class DatabaseManager:
    """
    Async database manager for CodeMCP storage layer.
    
    Handles database connections, sessions, and common CRUD operations
    with proper error handling and connection pooling.
    """
    
    def __init__(self, config: Config = None):
        self.config = config or Config()
        self.engine = None
        self.async_session_factory = None
        self.is_initialized = False
        
        # Database configuration
        self.database_url = self._get_database_url()
        self.is_sqlite = "sqlite" in self.database_url
        
        # Connection statistics
        self.stats = {
            "connections_created": 0,
            "queries_executed": 0,
            "transactions_committed": 0,
            "transactions_rolled_back": 0,
            "errors_encountered": 0
        }
        
        log_info(f"DatabaseManager initialized with URL: {self._sanitize_url(self.database_url)}")
    
    def _get_database_url(self) -> str:
        """Get database URL from configuration."""
        # Check environment variable first
        db_url = os.getenv("CODEMCP_DATABASE_URL")
        if db_url:
            return db_url
        
        # Use SQLite by default
        db_dir = Path.home() / ".codemcp" / "data"
        db_dir.mkdir(parents=True, exist_ok=True)
        db_path = db_dir / "codemcp.db"
        
        # For async SQLite, we need aiosqlite
        return f"sqlite+aiosqlite:///{db_path}"
    
    def _sanitize_url(self, url: str) -> str:
        """Sanitize database URL for logging (remove credentials)."""
        if "://" not in url:
            return url
        
        protocol, rest = url.split("://", 1)
        if "@" in rest:
            credentials, host_part = rest.rsplit("@", 1)
            return f"{protocol}://***:***@{host_part}"
        return url
    
    async def initialize(self):
        """Initialize database engine and create tables."""
        try:
            # Create async engine
            engine_kwargs = {
                "echo": self.config.debug,
                "future": True,
            }
            
            # SQLite-specific configuration
            if self.is_sqlite:
                engine_kwargs.update({
                    "poolclass": StaticPool,
                    "connect_args": {
                        "check_same_thread": False,
                        "timeout": 20
                    }
                })
            
            self.engine = create_async_engine(self.database_url, **engine_kwargs)
            
            # Create session factory
            self.async_session_factory = async_sessionmaker(
                self.engine,
                class_=AsyncSession,
                expire_on_commit=False
            )
            
            # Create tables
            async with self.engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
                log_info("Database tables created successfully")
            
            self.is_initialized = True
            log_info("Database manager initialized successfully")
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "database_initialization"})
            raise DatabaseError(f"Failed to initialize database: {error_response['message']}")
    
    async def close(self):
        """Close database connections."""
        if self.engine:
            await self.engine.dispose()
            log_info("Database connections closed")
    
    @asynccontextmanager
    async def get_session(self):
        """Get async database session with automatic cleanup."""
        if not self.is_initialized:
            await self.initialize()
        
        async with self.async_session_factory() as session:
            try:
                self.stats["connections_created"] += 1
                yield session
            except Exception as e:
                await session.rollback()
                self.stats["transactions_rolled_back"] += 1
                self.stats["errors_encountered"] += 1
                raise e
            finally:
                await session.close()
    
    async def execute_query(self, query: str, parameters: Dict[str, Any] = None) -> Any:
        """Execute a raw SQL query."""
        try:
            async with self.get_session() as session:
                result = await session.execute(text(query), parameters or {})
                await session.commit()
                self.stats["queries_executed"] += 1
                self.stats["transactions_committed"] += 1
                return result
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "execute_query",
                "query": query,
                "parameters": parameters
            })
            raise DatabaseError(f"Query execution failed: {error_response['message']}")
    
    async def create(self, model_instance: T) -> T:
        """Create a new record."""
        try:
            async with self.get_session() as session:
                session.add(model_instance)
                await session.commit()
                await session.refresh(model_instance)
                self.stats["transactions_committed"] += 1
                log_debug(f"Created {type(model_instance).__name__} with ID: {getattr(model_instance, 'id', 'N/A')}")
                return model_instance
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "create",
                "model": type(model_instance).__name__
            })
            raise DatabaseError(f"Failed to create record: {error_response['message']}")
    
    async def get_by_id(self, model_class: Type[T], record_id: str) -> Optional[T]:
        """Get a record by ID."""
        try:
            async with self.get_session() as session:
                result = await session.get(model_class, record_id)
                self.stats["queries_executed"] += 1
                return result
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "get_by_id",
                "model": model_class.__name__,
                "record_id": record_id
            })
            raise DatabaseError(f"Failed to get record: {error_response['message']}")
    
    async def get_by_filter(self, model_class: Type[T], **filters) -> List[T]:
        """Get records by filter criteria."""
        try:
            async with self.get_session() as session:
                query = session.query(model_class)
                for key, value in filters.items():
                    if hasattr(model_class, key):
                        query = query.filter(getattr(model_class, key) == value)
                
                result = await query.all()
                self.stats["queries_executed"] += 1
                return result
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "get_by_filter",
                "model": model_class.__name__,
                "filters": filters
            })
            raise DatabaseError(f"Failed to get records: {error_response['message']}")
    
    async def update(self, model_instance: T, **updates) -> T:
        """Update a record."""
        try:
            async with self.get_session() as session:
                # Merge the instance into the session
                merged_instance = await session.merge(model_instance)
                
                # Apply updates
                for key, value in updates.items():
                    if hasattr(merged_instance, key):
                        setattr(merged_instance, key, value)
                
                # Update timestamp if available
                if hasattr(merged_instance, 'updated_at'):
                    merged_instance.updated_at = datetime.utcnow()
                
                await session.commit()
                await session.refresh(merged_instance)
                self.stats["transactions_committed"] += 1
                
                log_debug(f"Updated {type(merged_instance).__name__} with ID: {getattr(merged_instance, 'id', 'N/A')}")
                return merged_instance
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "update",
                "model": type(model_instance).__name__,
                "updates": updates
            })
            raise DatabaseError(f"Failed to update record: {error_response['message']}")
    
    async def delete(self, model_instance: T) -> bool:
        """Delete a record."""
        try:
            async with self.get_session() as session:
                merged_instance = await session.merge(model_instance)
                await session.delete(merged_instance)
                await session.commit()
                self.stats["transactions_committed"] += 1
                
                log_debug(f"Deleted {type(model_instance).__name__} with ID: {getattr(model_instance, 'id', 'N/A')}")
                return True
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "delete",
                "model": type(model_instance).__name__
            })
            raise DatabaseError(f"Failed to delete record: {error_response['message']}")
    
    async def delete_by_id(self, model_class: Type[T], record_id: str) -> bool:
        """Delete a record by ID."""
        try:
            record = await self.get_by_id(model_class, record_id)
            if record:
                return await self.delete(record)
            return False
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "delete_by_id",
                "model": model_class.__name__,
                "record_id": record_id
            })
            raise DatabaseError(f"Failed to delete record: {error_response['message']}")
    
    async def count(self, model_class: Type[T], **filters) -> int:
        """Count records with optional filters."""
        try:
            async with self.get_session() as session:
                query = session.query(model_class)
                for key, value in filters.items():
                    if hasattr(model_class, key):
                        query = query.filter(getattr(model_class, key) == value)
                
                count = await query.count()
                self.stats["queries_executed"] += 1
                return count
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "count",
                "model": model_class.__name__,
                "filters": filters
            })
            raise DatabaseError(f"Failed to count records: {error_response['message']}")
    
    async def bulk_create(self, model_instances: List[T]) -> List[T]:
        """Create multiple records in a single transaction."""
        try:
            async with self.get_session() as session:
                session.add_all(model_instances)
                await session.commit()
                
                # Refresh all instances
                for instance in model_instances:
                    await session.refresh(instance)
                
                self.stats["transactions_committed"] += 1
                log_debug(f"Bulk created {len(model_instances)} records")
                return model_instances
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "bulk_create",
                "count": len(model_instances)
            })
            raise DatabaseError(f"Failed to bulk create records: {error_response['message']}")
    
    async def cleanup_expired_cache(self) -> int:
        """Clean up expired cache entries."""
        try:
            async with self.get_session() as session:
                now = datetime.utcnow()
                result = await session.execute(
                    text("DELETE FROM cache_entries WHERE expires_at < :now"),
                    {"now": now}
                )
                deleted_count = result.rowcount
                await session.commit()
                
                self.stats["transactions_committed"] += 1
                log_debug(f"Cleaned up {deleted_count} expired cache entries")
                return deleted_count
        except Exception as e:
            error_response = handle_error(e, {"operation": "cleanup_expired_cache"})
            raise DatabaseError(f"Failed to cleanup expired cache: {error_response['message']}")
    
    async def get_database_stats(self) -> Dict[str, Any]:
        """Get database statistics."""
        try:
            stats = {
                "connection_stats": self.stats.copy(),
                "database_url": self._sanitize_url(self.database_url),
                "is_initialized": self.is_initialized,
                "is_sqlite": self.is_sqlite
            }
            
            # Get table counts
            async with self.get_session() as session:
                for model_class in [CallTreeModel, CallTreeNodeModel, AnalysisResultModel, 
                                  FileChangeModel, CacheEntryModel]:
                    count = await session.query(model_class).count()
                    stats[f"{model_class.__tablename__}_count"] = count
            
            return stats
        except Exception as e:
            error_response = handle_error(e, {"operation": "get_database_stats"})
            raise DatabaseError(f"Failed to get database stats: {error_response['message']}")
    
    async def vacuum_database(self) -> bool:
        """Vacuum the database (SQLite only)."""
        if not self.is_sqlite:
            log_warning("Vacuum operation only supported for SQLite databases")
            return False
        
        try:
            await self.execute_query("VACUUM")
            log_info("Database vacuum completed")
            return True
        except Exception as e:
            error_response = handle_error(e, {"operation": "vacuum_database"})
            raise DatabaseError(f"Failed to vacuum database: {error_response['message']}")


# Global database manager instance
_global_database_manager = None


async def get_database_manager(config: Config = None) -> DatabaseManager:
    """Get the global database manager instance."""
    global _global_database_manager
    
    if _global_database_manager is None:
        _global_database_manager = DatabaseManager(config)
        await _global_database_manager.initialize()
    
    return _global_database_manager


async def close_database_manager():
    """Close the global database manager."""
    global _global_database_manager
    
    if _global_database_manager:
        await _global_database_manager.close()
        _global_database_manager = None